import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import torch
import yaml
import cv2
from model.SegMatch import SegMatch
from matchers.dual_softmax_matcher import DualSoftMaxMatcher
from dataset.mega_pose_est_mnn import MegaDepthPoseMNNBenchmark

def get_best_device(verbose = False):
    device = torch.device('cpu')
    if torch.cuda.is_available():
        device = torch.device('cuda')
    elif torch.backends.mps.is_available():
        device = torch.device('mps')
    else:
        device = torch.device('cpu')
    if verbose: print (f"Fastest device found is: {device}")
    return device


if __name__ == "__main__":
    with open("../config/inference.yaml", 'r') as f:
        config = yaml.load(f, Loader=yaml.FullLoader)
    device = get_best_device()
    detector = cv2.SIFT_create()
    descriptor = SegMatch(config)
    matcher = DualSoftMaxMatcher()
    # load weights
    weight_path = "/home/liyuke/lyk_work/segmatch/ckpts_desc/segmentation/001"
    print(weight_path)
    for name in ['backbone',"match_head"]:
        model_path = weight_path +'/{}.pth'.format(name)
        m = getattr(descriptor,name)
        model_param = torch.load(model_path, map_location='cuda')
        m.load_state_dict(model_param)
    descriptor.eval()
    # init dataset
    mega_1500 = MegaDepthPoseMNNBenchmark()
    mega_1500.benchmark(
        detector_model = detector,
        descriptor_model = descriptor,
        matcher_model = matcher)